#include "sub_custom_bf16.h"

using namespace AscendC;

class KernelSub {
public:
    __aicore__ inline KernelSub() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                uint32_t totalSizeX,
                                uint32_t totalSizeY,
                                uint32_t totalTileNum,
                                uint32_t tileLength)
    {
        this->totalSizeX = totalSizeX;
        this->totalSizeY = totalSizeY;
        this->totalTileNum = totalTileNum;
        this->tileLength = tileLength;
        this->typeSize = sizeof(half);

        xGm.GSetBuffer((__gm__ half*)x, totalSizeX);
        yGm.GSetBuffer((__gm__ half*)y, totalSizeY);
        zGm.GSetBuffer((__gm__ half*)z, totalSizeX);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * this->typeSize);
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * this->typeSize);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * this->typeSize);

    }
    __aicore__ inline void Process()
    {
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        uint32_t repNum = this->totalSizeX / this->totalSizeY;

        for (uint32_t rep = 0; rep < repNum; rep++) {
            this->repOffset = rep * this->totalSizeY;
            this->copyBytes = this->tileLength * this->typeSize;
            for (uint32_t i = block_id; i < this->totalTileNum; i += core_num) {
                if (i == this->totalTileNum - 1) { // last tile need pad
                    uint32_t lastTileLength = this->totalSizeY - (this->totalTileNum - 1) * this->tileLength;
                    this->copyBytes = lastTileLength * this->typeSize;
                }

                CopyIn(i);
                Compute();
                CopyOut(i);
            }
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<half> xLocal = inQueueX.AllocTensor<half>();
        LocalTensor<half> yLocal = inQueueY.AllocTensor<half>();
        DataCopyExtParams  copyParams = {1, this->copyBytes, 0, 0, 0};
        DataCopyPadExtParams<half> copyPadParams = {true, 0, 0, 0};
        DataCopyPad(xLocal, xGm[this->repOffset + progress * this->tileLength], copyParams, copyPadParams);
        DataCopyPad(yLocal, yGm[progress * this->tileLength], copyParams, copyPadParams);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    __aicore__ inline void Compute()
    {
        LocalTensor<half> xLocal = inQueueX.DeQue<half>();
        LocalTensor<half> yLocal = inQueueY.DeQue<half>();
        LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
        Sub(zLocal, xLocal, yLocal, this->tileLength);
        outQueueZ.EnQue<half>(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<half> zLocal = outQueueZ.DeQue<half>();
        DataCopyExtParams copyParams = {1, this->copyBytes, 0, 0, 0};
        DataCopyPad(zGm[this->repOffset + progress * this->tileLength], zLocal, copyParams);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    GlobalTensor<half> xGm;
    GlobalTensor<half> yGm;
    GlobalTensor<half> zGm;

    uint32_t totalSizeX;
    uint32_t totalSizeY;
    uint32_t totalTileNum;
    uint32_t tileLength;
    uint32_t copyBytes;
    uint32_t repOffset{0};
    uint32_t typeSize;
};

extern "C" __global__ __aicore__ void sub_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                                 uint32_t totalSizeX,
                                                 uint32_t totalSizeY,
                                                 uint32_t totalTileNum,
                                                 uint32_t tileLength,
                                                 uint32_t mode)
{
    if (mode == 0) {
        KernelSub op;
        op.Init(x, y, z, totalSizeX, totalSizeY, totalTileNum, tileLength);
        op.Process();
    } else {
#if defined(__CCE_KT_TEST__) || (__CCE_AICORE__ == 220)
        KernelSubBf16<bfloat16_t> op;
        op.Init(x, y, z, totalSizeX, totalSizeY, totalTileNum, tileLength);
        op.Process();
#endif
    }
}
